#!/usr/bin/env python3
"""
测试VisualLocalizer的脚本
"""

import argparse
import asyncio
import time
from pathlib import Path

import cv2

from visual_localization import VisualLocalizer


def main():
    parser = argparse.ArgumentParser(
        description="测试VisualLocalizer定位功能",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
  python test_visual_localizer.py \\
    --image_path query.jpg \\
    --reconstruction_path /path/to/sfm_model \\
    --db_global_features /path/to/global-feats-netvlad.h5 \\
    --db_local_features /path/to/feats-superpoint-n4096-r1024.h5 \\
    --top_k 20
        """
    )
    parser.add_argument("--image_path", type=str, required=True,
                       help="查询图像路径")
    parser.add_argument("--reconstruction_path", type=str, required=True,
                       help="SfM重建模型路径")
    parser.add_argument("--db_global_features", type=str, required=True,
                       help="数据库全局特征文件路径 (NetVLAD, .h5)")
    parser.add_argument("--db_local_features", type=str, required=True,
                       help="数据库局部特征文件路径 (SuperPoint, .h5)")
    parser.add_argument("--top_k", type=int, default=20,
                       help="检索图像数量 (默认: 20)")
    args = parser.parse_args()

    print("=" * 80)
    print("VisualLocalizer 测试")
    print("=" * 80)
    print(f"查询图像: {args.image_path}")
    print(f"SfM模型: {args.reconstruction_path}")
    print(f"全局特征: {args.db_global_features}")
    print(f"局部特征: {args.db_local_features}")
    print(f"检索数量: {args.top_k}")
    print("=" * 80)

    # 验证文件存在
    for path_name, path_value in [
        ("查询图像", args.image_path),
        ("SfM模型", args.reconstruction_path),
        ("全局特征", args.db_global_features),
        ("局部特征", args.db_local_features),
    ]:
        if not Path(path_value).exists():
            raise FileNotFoundError(f"{path_name}不存在: {path_value}")

    # 读取图像
    print(f"\n[1/3] 读取查询图像...")
    img = cv2.imread(args.image_path, cv2.IMREAD_COLOR)
    if img is None:
        raise FileNotFoundError(f"无法读取图像: {args.image_path}")
    print(f"  图像尺寸: {img.shape[1]} x {img.shape[0]} (宽x高)")
    print(f"  图像通道: {img.shape[2]}")

    # 初始化VisualLocalizer
    print(f"\n[2/3] 初始化VisualLocalizer...")
    localizer = VisualLocalizer()

    t0 = time.time()
    ok = localizer.initialize(
        reconstruction_path=args.reconstruction_path,
        db_global_features_path=args.db_global_features,
        db_local_features_path=args.db_local_features,
        top_k=args.top_k,
        params={}
    )
    t_init = time.time() - t0
    if not ok:
        raise RuntimeError("初始化VisualLocalizer失败")

    print(f"✓ 初始化完成 (耗时: {t_init:.2f}s)")

    # 执行定位
    print(f"\n[3/3] 执行定位...")
    t0 = time.time()
    result = asyncio.run(localizer.localize(img))
    t_localize = time.time() - t0

    # 打印结果
    print("\n" + "=" * 80)
    print("定位结果")
    print("=" * 80)
    print(f"成功: {result.success}")
    
    if result.success:
        print(f"\n位姿信息:")
        print(f"  平移向量 (x, y, z): {result.translation}")
        print(f"  四元数 (w, x, y, z): {result.quaternion}")
        print(f"\n匹配统计:")
        print(f"  内点数量: {result.inlier_count}")
        if result.additional_info:
            print(f"  对应点数量: {result.additional_info.get('num_correspondences', 'N/A')}")
            retrieved = result.additional_info.get('retrieved', [])
            print(f"  检索到的图像: {len(retrieved)} 张")
            if len(retrieved) > 0:
                print(f"    前3张: {retrieved[:3]}")
        print(f"\n误差信息:")
        print(f"  重投影误差: {result.reprojection_error:.4f}")
        print(f"  置信度: {result.confidence:.4f}")
    else:
        print("定位失败")
        if result.additional_info:
            print(f"  信息: {result.additional_info}")
    
    print(f"\n时间统计:")
    print(f"  初始化时间: {t_init:.3f}s")
    print(f"  定位时间: {t_localize:.3f}s")
    print(f"  总时间: {t_init + t_localize:.3f}s")
    print("=" * 80)


if __name__ == "__main__":
    main()


